/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: 91093 ReduceScatter小数据量MTE实现，需Launch 16个aiv
 */
 
#include "aiv_communication_base.h"
 
using namespace AscendC;
 
class AivReduceScatterSmallGraph91093 : public AivCommBase {
public:
    __aicore__ inline AivReduceScatterSmallGraph91093() {}
 
    template<typename T>
    __aicore__ inline void Process(GM_ADDR input, GM_ADDR output, uint64_t len, int32_t tag);
};

template <typename T>
__aicore__ inline void AivReduceScatterSmallGraph91093::Process(GM_ADDR input, GM_ADDR output, uint64_t len,
    int32_t tag)
{
    localSetTensor.SetValue(0, tag);
    __gm__ T *inputGM = (__gm__ T *)input;
    __gm__ T *outputGM = (__gm__ T *)output;
 
    uint32_t blockNumPerGroup = block_num / rankSize_; // block_num需要能被rankSize_整除
    uint32_t blockIdxInGroup = block_idx % blockNumPerGroup;

    uint32_t padCount = UB_ALIGN_SIZE / sizeof(T);
    uint64_t avgLengthPerBlock = CeilDiv(len, blockNumPerGroup);
    uint64_t avgLengthPerSlice = CeilDiv(avgLengthPerBlock, padCount) * padCount; // 32B对齐
    uint64_t sliceCount = CeilDiv(len, avgLengthPerSlice);
    uint64_t tailLength = len - (sliceCount - 1) * avgLengthPerSlice;

    uint64_t count = CalActualCount(blockIdxInGroup, sliceCount, avgLengthPerSlice, tailLength);
    uint64_t blockOffset = blockIdxInGroup * avgLengthPerSlice;
    uint32_t dstRank = block_idx / blockNumPerGroup;

    uint32_t flagOffsetBase = BASE_FLAG_OFFSET * AIV_REDUCE_SCATTER_91093_SMALLDATA_GRAPH;
    uint32_t flagXOffset = blockIdxInGroup * FLAG_SIZE + rank_ * blockNumPerGroup * FLAG_SIZE + flagOffsetBase;
    uint32_t flagOffset = block_idx * FLAG_SIZE + flagOffsetBase;

    __gm__ int32_t *ctrlFlagsGMX = (__gm__ int32_t *)(GM_OUT[dstRank] + flagXOffset);
    __gm__ int32_t *ctrlFlagsGM = (__gm__ int32_t *)(GM_OUT[rank_] + flagOffset);
    __gm__ int32_t *ctrlFlagsGML = (__gm__ int32_t *)(GM_OUT[rank_] + flagXOffset);
    GlobalTensor<int32_t> globalSet;
    globalSet.SetGlobalBuffer(ctrlFlagsGML, UB_FLAG_PAD_COUNT);
    
    if (dstRank == rank_) {
        CpGM2GM(outputGM + blockOffset, (__gm__ T *)(inputGM + rank_ * len + blockOffset), count);
        pipe_barrier(PIPE_MTE3);
        DataCopy(globalSet, localSetTensor, UB_FLAG_PAD_COUNT);
    } else {
        globalSet.SetGlobalBuffer(ctrlFlagsGMX, UB_FLAG_PAD_COUNT);
        DataCopy(globalSet, localSetTensor, UB_FLAG_PAD_COUNT);
        CheckFlagNew(ctrlFlagsGM, tag); // 跨片
        CheckFlagNew(ctrlFlagsGML, tag); // 本地
        PipeBarrier<PIPE_ALL>();

        CpGM2GM(outputGM + blockOffset, (__gm__ T *)(GM_IN[dstRank]) + rank_ * len + blockOffset, count, true,
            reduceOp_);

        ctrlFlagsGMX = (__gm__ int32_t *)(GM_OUT[dstRank] + block_num * FLAG_SIZE + flagXOffset);
        ctrlFlagsGM = (__gm__ int32_t *)(GM_OUT[rank_] + block_num * FLAG_SIZE + flagOffset);
        pipe_barrier(PIPE_MTE3);
        globalSet.SetGlobalBuffer(ctrlFlagsGMX, UB_FLAG_PAD_COUNT);
        DataCopy(globalSet, localSetTensor, UB_FLAG_PAD_COUNT);
        CheckFlagNew(ctrlFlagsGM, tag);       
    }

    return;
}

template<typename T>
__aicore__ inline void aiv_reduce_scatter_91093_smalldata_graph(KERNEL_ARGS_DEF)
{
    AivReduceScatterSmallGraph91093 op;
    op.Init(KERNEL_CLASS_INIT, true);
    op.Process<T>(input, output, len, tag);
}